import uuid
from typing import Optional, Callable, Any
from datetime import timedelta
from utils.logger import Logger
from utils.exceptions import LockAcquireError
from middleware.redis_adapter import get_redis_adapter
from config import middleware_config

# 初始化日志
logger = Logger.get_logger(agent_name="distributed_lock")


class DistributedLock:
    """基于 Redis 的分布式锁（支持超时、重试、原子释放）"""

    def __init__(self):
        self.redis_client = get_redis_adapter()  # 使用Redis适配器，支持集群模式
        self.lock_prefix: str = "distributed:lock:"  # 锁键前缀
        # 从配置读取默认参数（支持环境变量覆盖）
        self.default_timeout: int = middleware_config.redis.lock_timeout  # 锁超时时间（秒）
        self.default_retries: int = middleware_config.redis.lock_retries  # 重试次数
        self.default_retry_delay: float = middleware_config.redis.lock_retry_delay  # 重试间隔（秒）
        logger.info("分布式锁初始化完成，默认超时: {}s".format(self.default_timeout))

    def _get_lock_key(self, lock_name: str) -> str:
        """生成锁的 Redis 键名"""
        return f"{self.lock_prefix}{lock_name}"

    def acquire(
            self,
            lock_name: str,
            timeout: int = None,
            retries: int = None,
            retry_delay: float = None
    ) -> Optional[str]:
        """
        获取分布式锁
        Args:
            lock_name: 锁名称（唯一标识资源）
            timeout: 锁超时时间（秒），默认使用配置值
            retries: 获取失败重试次数，默认使用配置值
            retry_delay: 重试间隔（秒），默认使用配置值
        Returns:
            锁标识（解锁时需传入，防止误解锁），失败返回 None
        """
        timeout = timeout or self.default_timeout
        retries = retries or self.default_retries
        retry_delay = retry_delay or self.default_retry_delay
        lock_key = self._get_lock_key(lock_name)
        lock_id = str(uuid.uuid4())  # 唯一锁标识
        retry_count = 0

        while retry_count <= retries:
            try:
                # Redis SET NX EX 原子操作：不存在则设置，同时设置过期时间
                acquired = self.redis_client.set(
                    key=lock_key,
                    value=lock_id,
                    nx=True,  # 仅当键不存在时成功
                    ex=timeout  # 自动过期，避免死锁
                )

                if acquired:
                    logger.debug(f"获取锁成功: lock_name={lock_name}, lock_id={lock_id[:8]}")
                    return lock_id
                else:
                    # 锁已被占用，重试
                    retry_count += 1
                    logger.debug(f"获取锁失败（已被占用）: lock_name={lock_name}, 重试={retry_count}/{retries}")
                    if retry_count <= retries:
                        import time
                        time.sleep(retry_delay)
            except Exception as e:
                logger.error(f"获取锁异常: lock_name={lock_name}, error={str(e)}", exc_info=True)
                retry_count += 1
                if retry_count <= retries:
                    import time
                    time.sleep(retry_delay)

        logger.error(f"获取锁失败（重试耗尽）: lock_name={lock_name}")
        return None

    def acquire_with_exception(
            self,
            lock_name: str,
            timeout: int = None,
            retries: int = None,
            retry_delay: float = None
    ) -> str:
        """获取锁（失败抛出异常）"""
        lock_id = self.acquire(lock_name, timeout, retries, retry_delay)
        if not lock_id:
            raise LockAcquireError(
                message=f"分布式锁获取失败（重试{retries or self.default_retries}次）",
                context={"lock_name": lock_name}
            )
        return lock_id

    def release(self, lock_name: str, lock_id: str) -> bool:
        """
        释放分布式锁（仅释放自己获取的锁）
        Args:
            lock_name: 锁名称
            lock_id: 锁标识（acquire 返回的值）
        Returns:
            释放成功返回 True，失败返回 False
        """
        lock_key = self._get_lock_key(lock_name)
        try:
            # Lua 脚本保证原子性：判断锁标识匹配才删除
            lua_script = """
            if redis.call('get', KEYS[1]) == ARGV[1] then
                return redis.call('del', KEYS[1])
            else
                return 0
            end
            """
            # 替换 ARGV[1] 为实际 lock_id（避免脚本注入风险）
            lua_script = lua_script.replace("ARGV[1]", f"'{lock_id}'")
            result = self.redis_client.eval(lua_script, keys=[lock_key], args=[])

            if result == 1:
                logger.debug(f"释放锁成功: lock_name={lock_name}, lock_id={lock_id[:8]}")
                return True
            else:
                logger.warning(f"释放锁失败: lock_name={lock_name}, lock_id={lock_id[:8]}（锁不存在或标识不匹配）")
                return False
        except Exception as e:
            logger.error(f"释放锁异常: lock_name={lock_name}, lock_id={lock_id[:8]}, error={str(e)}", exc_info=True)
            return False

    def with_lock(
            self,
            lock_name: str,
            timeout: int = None,
            retries: int = None,
            retry_delay: float = None
    ) -> Callable:
        """装饰器：自动获取/释放锁"""

        def decorator(func: Callable) -> Callable:
            def wrapper(*args, **kwargs) -> Any:
                lock_id = None
                try:
                    lock_id = self.acquire_with_exception(lock_name, timeout, retries, retry_delay)
                    return func(*args, **kwargs)
                finally:
                    if lock_id:
                        self.release(lock_name, lock_id)

            return wrapper

        return decorator


# 单例实例（全局复用）
distributed_lock = DistributedLock()

# 测试入口
if __name__ == "__main__":
    test_logger = Logger.update_context(task_id="distributed_lock_test")

    # 测试基础获取/释放
    lock_name = "test_resource"
    lock_id = distributed_lock.acquire(lock_name, timeout=5)
    if lock_id:
        distributed_lock.release(lock_name, lock_id)


    # 测试装饰器
    @distributed_lock.with_lock("test_decorator_lock", timeout=3)
    def test_func():
        test_logger.info("装饰器模式：获取锁后执行业务")


    test_func()

    # 测试并发竞争（模拟 2 个线程抢锁）
    import threading


    def compete_lock(thread_id):
        try:
            lock_id = distributed_lock.acquire_with_exception(f"compete_lock", retries=2)
            test_logger.info(f"线程{thread_id}：抢锁成功")
            import time
            time.sleep(1)
        except LockAcquireError:
            test_logger.error(f"线程{thread_id}：抢锁失败")
        finally:
            if lock_id:
                distributed_lock.release(f"compete_lock", lock_id)


    threads = [threading.Thread(target=compete_lock, args=(i,)) for i in range(2)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()